//
//  protocal.c
//  myvpn
//
//  Created by Zhuobin He on 12-12-31.
//  Copyright (c) 2012年 Zhuobin He. All rights reserved.
//


#include "SysLib.h"
#include "Firewall.h"
#include "Protocal.h"
#include "SocketUtil.h"
#include "Tunnel.h"

#include "Common.h"
#include "Md5.h"
#include "TunDev.h"

int prime_num[] = {3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97};
int prime_count = sizeof(prime_num) / sizeof(3);

int parse_int16(char *buffer) {
    uint16_t plength;
    memcpy((void *)&plength, buffer, sizeof(plength));
    return (int16_t)ntohs(plength);

}

int decode_data(PBuffer read_buffer, char *tun_buffer) {
    int readable_size = buffer_readable(read_buffer);
    if (readable_size < 6) return -1;
    
    char *pread = buffer_pread(read_buffer);
    // Get length mark
    int seed = parse_int16(pread);
    int seed2 = parse_int16(pread + 2);
    int ret = seed2 + seed;
    int plen = ret;
    if (plen < 0) plen = 0;

    if (plen > 1500) {
        log_error("Get illegal length mark %d", plen);
        return PTYPE_ILLEGAL;
    }

    // Decode data
    {
        int padding = 0;
        if ((plen + 4) % 16) {
            padding = 16 - (plen + 4) % 16;
        }
        
        if (readable_size < plen + 4 + padding) {
            log_trace("Decode: not enough data to decode, get only %d, need %d", readable_size, plen + 4 + padding);
            return PTYPE_PARTIAL;
        }

        pread += 4;
        buffer_pread_inc(read_buffer, plen + 4 + padding);
        
        int i;
        for (i=0; i<plen; i++) {
            int tmp = (pread[i] - i * seed - seed2) % 256;
            if (tmp < 0) tmp += 256;
            tun_buffer[i] = (char)tmp;
        }
    }
    return ret;
}

void put_int16(char *buffer, int value) {
    uint16_t plength = htons((uint16_t)((int16_t)value));
    memcpy(buffer, (const char *)&plength, 2);
}


int encode_data(char *tun_buffer, int size, PBuffer write_buffer) {
    int encoded_size = 4;
    if (size > 0) encoded_size += size;
    int padding = 0;
    if (encoded_size % 16) {
        padding = 16 - encoded_size % 16;
        encoded_size += padding;
    }
    if (buffer_writable(write_buffer) < encoded_size) {
        log_info("Write buffer is full, drop tun package");
        return 0;
    }

    char* pwrite = buffer_pwrite(write_buffer);

    int seed = prime_num[random() % prime_count];
    int seed2 = size - seed;
    put_int16(pwrite, seed);
    put_int16(pwrite + 2, seed2);
    pwrite += 4;

    // Encode data
    if (size > 0) {
        int i;
        for (i=0; i<size; i++) {
            pwrite[i] = (char)((tun_buffer[i] + i * seed + seed2) % 256);
        }
        pwrite += size;
    }
    // padding
    {
        int i;
        for (i=0; i<padding; i++) {
            pwrite[i] = (char) random();
        }
    }
    buffer_pwrite_inc(write_buffer, encoded_size);
    return encoded_size;
}


void md5(char *src, char *hex_output) {
    md5_state_t state;
    md5_byte_t digest[16];
    int di;
    md5_init(&state);
    md5_append(&state, (const md5_byte_t *)src, (int)strlen(src));
    md5_finish(&state, digest);
    for (di = 0; di < 16; ++di)
        sprintf(hex_output + di * 2, "%02x", digest[di]);
}

bool is_http_header_finish(PBuffer read_buf) {
    int size = buffer_readable(read_buf);
    char *pread = buffer_pread(read_buf);
    pread[size] = 0;
    return strstr(pread, "\n\r\n") || strstr(pread, "\n\n");
}

struct _Header {
    char *action;
    char **key;
    char **value;
    int count;
};
typedef struct _Header Header;
typedef Header * PHeader;

char * get_sub_str(char *src, char *tail) {
    int len = (int)(tail - src);
    char *dst = (char *) malloc(len + 1);
    memcpy(dst, src, sizeof(char) * len);
    dst[len] = '\0';
    return dst;
}

PHeader header_buffer = NULL;

void dispose_protocal_buffer() {
    if (header_buffer) {
        int i;
        for (i=0; i<10; i++) {
            if (header_buffer->key && header_buffer->key[i]) free(header_buffer->key[i]);
            if (header_buffer->value && header_buffer->value[i]) free(header_buffer->value[i]);
        }
        if (header_buffer->action) free(header_buffer->action);
        if (header_buffer->key) free(header_buffer->key);
        if (header_buffer->value) free(header_buffer->value);
        
        header_buffer = NULL;
    }
}

PHeader parse_header(PBuffer buf) {
    dispose_protocal_buffer();
    int head_size = 0;
    char *head = buffer_pread(buf);
    char *iter = strstr(head, "\n\r\n");
    if (iter) {
        head_size = iter - head;
    }
    iter = strstr(head, "\n\n");
    if (iter && iter - head < head_size) {
        head_size = iter - head;
    } else {
        iter = NULL;
    }
    head[head_size] = '\n';
    head[head_size + 1] = '\0';
    head_size += iter ? 2 : 3;
    
    iter = strchr(head, '\n');
    
    PHeader ret = NULL;
    zmalloc(&ret, sizeof(Header));
    header_buffer = ret;
    zmalloc(&ret->key, sizeof(char *) * 10);
    zmalloc(&ret->value, sizeof(char *) * 10);
    ret->action = trim(get_sub_str(head, iter));

    head = iter + 1;
    int i;
    for (i=0; i<10; i++) {
        iter = strchr(head, ':');
        if (!iter) break;
        ret->key[i] = get_sub_str(head, iter);
        head = iter;
        for (;head[0] != 0 && head[0] != ' '; head++) ;
        iter = strchr(head, '\n');
        if (!iter) break;
        ret->value[i] = trim(get_sub_str(head, iter));
        head = iter + 1;
        ret->count ++;
    }
    buffer_pread_inc(buf, head_size);
    buffer_compact(buf);
    return ret;
}

int find_header(PHeader header, const char *key) {
    int i;
    for (i=0; i<header->count; i++) {
        if (strcmp(header->key[i], key) == 0) return i;
    }
    return -1;
}

bool _client_auth(PTunnel tunnel, PVpnConfig conf) {
    if (tunnel->status == TSTATUS_AUTH_BEGIN) {
        char *auth_buf = buffer_pwrite(tunnel->write_buffer);
        strcpy(auth_buf, "GET /auth HTTP/1.1\r\n");
        strcat(auth_buf, "Host: lab.qq.com\r\nConnection: keep-alive\r\n");
        sprintf(auth_buf + strlen(auth_buf), "Username: %s\r\n", tunnel->username);
        sprintf(auth_buf + strlen(auth_buf), "Seed: %d\r\n", tunnel->client_auth_seed);
        strcat(auth_buf, "User-Agent: generic\r\n\r\n");
        log_debug("Sending auth request: %s", auth_buf);
        buffer_pwrite_inc(tunnel->write_buffer, (int)strlen(auth_buf));
        tunnel->status += 1;
    } else if (is_http_header_finish(tunnel->read_buffer)) {
        log_debug("Read auth response:");
        log_debug(buffer_pread(tunnel->read_buffer));
        PHeader header = parse_header(tunnel->read_buffer);
        switch (tunnel->status) {
            case TSTATUS_AUTH_BEGIN + 1:
            {
                log_info("Recv auth response");
                if (strcmp(header->action, "HTTP/1.1 401 Unauthorized")) {
                    return false;
                }
                
                int iter;
                if ((iter = find_header(header, "Seed")) == -1) {
                    return false;
                }
                char tmp[100];
                strcpy(tmp, header->value[iter]);
                tunnel->auth_seed = str2int(header->value[iter]);
                strcat(tmp, conf->users[0]->password);
                char digest[40];
                md5(tmp, digest);
                
                char *auth_buf = buffer_pwrite(tunnel->write_buffer);
                strcpy(auth_buf, "GET /download HTTP/1.1\r\n");
                strcat(auth_buf, "Host: lab.qq.com\r\nConnection: keep-alive\r\n");
                strcat(auth_buf, "User-Agent: generic\r\n");
                sprintf(auth_buf + strlen(auth_buf), "RouteConfVer: %d\r\n", conf->route_conf_ver);
                strcat(auth_buf, "Digest: ");
                strcat(auth_buf, digest);
                strcat(auth_buf, "\r\n\r\n");
                buffer_pwrite_inc(tunnel->write_buffer, (int)strlen(auth_buf));
                tunnel->status += 1;
                log_debug("Sending download");
                log_debug(buffer_pread(tunnel->write_buffer));
                break;
            }
            case TSTATUS_AUTH_BEGIN + 2:
            {
                log_info("Recv download response");
                if (strcmp(header->action, "HTTP/1.1 200 OK")) {
                    return false;
                }

                int iter;
                if ((iter = find_header(header, "ClientIp")) == -1) {
                    return false;
                }
                char *client_ip = header->value[iter];
                
                if ((iter = find_header(header, "ServerIp")) == -1) {
                    return false;
                }
                char *server_ip = header->value[iter];
                tunnel_establish(tunnel, client_ip, server_ip);
                if ((iter = find_header(header, "PushRouteVer")) != -1) {
                    int route_conf_ver = str2int(header->value[iter]);
                    if ((iter = find_header(header, "PushRoutes")) != -1) {
                        char *s1 = header->value[iter];
                        FILE *rfile = fopen(conf->route_cache_file, "w");
                        char tmp[50] = "";
                        sprintf(tmp, "#VER:%d\r\n", route_conf_ver);
                        int wret = fputs(tmp, rfile);
                        wret = (wret >= 0) ? fputs("#Auto generated by program. Do not modify\r\n#Modify [routes] section in config file instead if you need custom routes\r\n", rfile) : -1;
                        for (;wret >= 0;) {
                            char *s2 = strchr(s1, ',');
                            if (s2 == NULL) break;
                            s2[0] = '\0';
                            sprintf(tmp, "%s\r\n", s1);
                            wret = (wret >= 0) ? fputs(tmp, rfile) : -1;
                            s1 = s2 + 1;
                        }

                        if (!rfile) {
                            log_error("Error openning push route file for write");
                            return false;
                        }


                        fclose(rfile);
                        if (wret < 0) {
                            log_error("Error writing push route file");
                            return false;
                        }
                    }
                }

            }
        }
    }
    return true;
}

void client_auth(PTunnel tunnel, PVpnConfig conf) {
    if (!_client_auth(tunnel, conf)) {
        tunnel_close(tunnel, true);
    }
    dispose_protocal_buffer();
}

bool gen_close_response(Tunnel *tunnel) {
    char *auth_buf = buffer_pwrite(tunnel->write_buffer);
    strcpy(auth_buf, "HTTP/1.1 500 Unauthorized\r\n");
    strcat(auth_buf, "Server: Apache\r\n");
    strcat(auth_buf, "Connection: close\r\n");
    strcat(auth_buf, "Content-length: 16\r\n\r\n");
    strcat(auth_buf, "500 Unauthorized");
    buffer_pwrite_inc(tunnel->write_buffer, (int)strlen(auth_buf));
    tunnel_close(tunnel, false);
    log_error("Invalid client auth");
    return false;
}



bool assign_ip(PTunnel tunnel, PVpnConfig config) {
    log_debug("Assigning ip");
    int ip_id = -1;
    int i;
    int max_match = 0;
    for (i=3; i<256; i++) {
        int curr_match = 0;
        PTunnelMeta meta = tunnel_get_meta_by_ip_id(i);
        if (strcmp(tunnel->username, meta->username) == 0
                && tunnel->client_auth_seed == meta->auth_seed) {
            curr_match = 3;
        } else if (strcmp(tunnel->username, meta->username) == 0
                && !meta->connected) {
            curr_match = 2;
        } else if (!meta->connected) {
            curr_match = 1;
        }
        if (curr_match > max_match) {
            max_match = curr_match;
            ip_id = i;
        }
    }
    if (max_match) {
        PTunnel old_tunnel = tunnel_get_by_ip_id(ip_id);
        if (old_tunnel) {
            old_tunnel->duplicated = true;
            tunnel_close(old_tunnel, false);
        }

        log_info("Assign ip: %d, match: %d", ip_id, max_match);
        char client_ip[20];
        char server_ip[20];
        sprintf(client_ip, "%s%d", config->network, ip_id);
        sprintf(server_ip, "%s%d", config->network, 1);
        tunnel_establish(tunnel, client_ip, server_ip);
        return true;
    } else {
        log_error("No ip available, username: %s", tunnel->username);
    }
    return false;
}

bool _server_auth(PTunnel tunnel, PVpnConfig conf) {
    PUserInfo *users = conf->users;

    if (is_http_header_finish(tunnel->read_buffer)) {
        log_debug("Receive auth data: %s", buffer_pread(tunnel->read_buffer));
        PHeader header = parse_header(tunnel->read_buffer);
        
        switch (tunnel->status - TSTATUS_AUTH_BEGIN) {
            case 0:
            {
                if (strcmp(header->action, "GET /auth HTTP/1.1")) {
                    return false;
                }
                int iter;

                if ((iter = find_header(header, "Username")) == -1) return false;
                char *username = header->value[iter];
                tunnel->username = string_alloc(username);
                log_info("Recv auth, user: %s", tunnel->username);

                if ((iter = find_header(header, "Seed")) == -1) return false;
                tunnel->client_auth_seed = str2int(header->value[iter]);
                log_debug("Client auth seed: %d", tunnel->client_auth_seed);

                char *auth_buf = buffer_pwrite(tunnel->write_buffer);
                strcpy(auth_buf, "HTTP/1.1 401 Unauthorized\r\n");
                sprintf(auth_buf + strlen(auth_buf), "Seed: %d\r\n", tunnel->auth_seed = (int)(random() % 30000));
                strcat(auth_buf, "Server: Apache\r\nConnection: keep-alive\r\nContent-Type: application/octet-stream\r\nContent-Length: 0\r\n\r\n");
                log_debug("Server auth seed: %d", tunnel->auth_seed);
                buffer_pwrite_inc(tunnel->write_buffer, (int)strlen(auth_buf));
                log_debug("Sending auth resp");
                log_debug(auth_buf);
                tunnel->status += 1;
                
                break;
            }
            case 1:
            {
                if (strcmp(header->action, "GET /download HTTP/1.1")) {
                    return false;
                }
                log_info("Recv download");
                
                int iter;
                char server_digest[40] = "";
                int i;
                for (i=0; users[i] != NULL; i++) {
                    if (strcmp(tunnel->username, users[i]->username) == 0) {
                        char tmp[100];
                        sprintf(tmp, "%d%s", tunnel->auth_seed, users[i]->password);
                        md5(tmp, server_digest);
                        break;
                    }
                }
                if (server_digest[0] == '\0') {
                    return false;
                }
                
                if ((iter = find_header(header, "Digest")) == -1) return false;
                char *client_digest = header->value[iter];
                
                if (strcmp(server_digest, client_digest) != 0) {
                    log_error("Digest not match, act: %s, exp: %s", client_digest, server_digest);
                    return false;
                }
                log_debug("Digest match, username: %s", tunnel->username);
                if ((iter = find_header(header, "Heartbeat")) != -1) {
                    char *pp = strchr(header->value[iter], ' ');
                    if (pp) {
                        pp[0] = 0;
                        int heartbeat = str2int(header->value[iter]);
                        int timeout = str2int(pp + 1);
                        if (heartbeat >= 15 && heartbeat <= 360 && timeout >= 30 && timeout <= 600) {
                            tunnel->heartbeat = heartbeat;
                            tunnel->timeout = timeout;
                            log_info("Custom heartbeat: %d, timeout: %d", heartbeat, timeout);
                        }
                    }
                }

                if (!assign_ip(tunnel, conf)) {
                    return false;
                }

                char *auth_buf = buffer_pwrite(tunnel->write_buffer);
                strcpy(auth_buf, "HTTP/1.1 200 OK\r\n");
                strcat(auth_buf, "ClientIp: ");
                strcat(auth_buf, tunnel->client_ip);
                strcat(auth_buf, "\r\nServerIp: ");
                strcat(auth_buf, tunnel->server_ip);
                if ((iter = find_header(header, "RouteConfVer")) != -1 && conf->routes) {
                    int route_conf_ver = str2int(header->value[iter]);
                    if (route_conf_ver < conf->route_conf_ver) {
                        char *pbuf = auth_buf + strlen(auth_buf);
                        sprintf(pbuf, "\r\nPushRouteVer: %d\r\nPushRoutes: ", conf->route_conf_ver);
                        PRoute *route;
                        for (route = conf->routes; *route; route++) {
                            pbuf += strlen(pbuf);
                            sprintf(pbuf, "%s/%d,", route[0]->address, route[0]->prefix);
                        }
                    }
                }
                strcat(auth_buf, "\r\nServer: Apache\r\nConnection: keep-alive\r\nContent-Type: application/octet-stream\r\nTransfer-Encoding: chunked\r\n\r\n");
                buffer_pwrite_inc(tunnel->write_buffer, (int)strlen(auth_buf));
                log_debug("Sending download resp");
                log_debug(auth_buf);
            }
        }
    } else {
        log_debug("request data not finish");
    }
    return true;
}

void server_auth(PTunnel tunnel, PVpnConfig conf) {
    if (!_server_auth(tunnel, conf)) {
        gen_close_response(tunnel);
    }
    dispose_protocal_buffer();
}

